from argparse import ArgumentParser
from pathlib import Path
from loguru import logger
from distil_whisper import WhisperPipe
from audio_slicer import get_audio_segments, get_audio_slice_tmp_file
import pandas as pd
from tqdm import tqdm

from script_gen import ScriptGenerator

def get_distance_from_timestamp(stage_details: pd.DataFrame, timestamp: str):
    # the timestamp could be not in the stage_details, so we need to find the closest timestamp
    # get the closest timestamp
    closest_timestamp = stage_details.iloc[(stage_details['timestamp'] - timestamp).abs().argsort()[:1]]
    # get the distance
    distance = closest_timestamp['distance'].values[0]
    return distance

def main():
    parser = ArgumentParser()
    # /home/andy/lab/distil-large-v2-local
    parser.add_argument("--model_id", type=str, help="model id on huggingface", required=True)
    # /mnt/c/Users/andy/Downloads/default/瑞典_Eksharad
    parser.add_argument("--track_path", type=str, help="path to track including audio files and the stage_details.csv", required=False)
    parser.add_argument("--tracks_path", type=str, help="path to include all the tracks folders")
    parser.add_argument("--output_path", type=str, help="path to output folder, default to the same folder of track_path", required=False)
    parser.add_argument("--noise_reduction", action="store_true", help="noise reduction before recognization")
    args = parser.parse_args()
    args = vars(args)


    logger.info(args)

    # load the model
    pipe = WhisperPipe(args)
    
    # loop through all the folders
    tracks_paths = []
    if not args['tracks_path']:
        tracks_paths.append(Path(args['track_path']))
    else:
        tracks_paths = [p for p in Path(args['tracks_path']).rglob('*') if p.is_dir()]

    for track_path in tqdm(tracks_paths):
        # trackname is the folder name of the track_path
        trackname = Path(track_path)

        # load the stage details including timestamp,stage_time,distance
        stage_details = pd.read_csv(track_path / 'stage_details.csv')

        # get all the wav files in the track
        wav_files = list(Path(track_path).rglob('*.wav'))

        # loop through all the wav files
        logger.info(f'going to loop through all the wav files in the track {track_path}')
        lines_to_write = []
        for wav_file in wav_files:
            # the file name is distance-stage_time-timestamp.wav
            file_name = wav_file.name
            file_name_parts = file_name.split('-')
            has_minus_distance = len(file_name_parts) == 4
            file_names_start_index = 1 if has_minus_distance else 0
            distance = 0 - int(file_name_parts[file_names_start_index]) if has_minus_distance else file_name_parts[file_names_start_index]
            stage_time = file_name_parts[file_names_start_index+1]
            timestamp = file_name_parts[file_names_start_index+2].split('.')[0]    # 10,000,000 ticks per second

            # split the audio into segments
            noise_reduced_audio, segments = get_audio_segments(wav_file, args['noise_reduction'])

            # loop through all the segments
            for segment in segments:
                # get the slice tmp file
                tmp_file = get_audio_slice_tmp_file(noise_reduced_audio, segment)

                # recognize the slice
                result = pipe.recognize(tmp_file)

                # get the timestamp appended with the segment start time in seconds
                timestamp_with_segment_start_time = int(timestamp) + segment[0] * 10000000

                # get the distance from the timestamp
                the_distance = get_distance_from_timestamp(stage_details, timestamp_with_segment_start_time)

                # save the result to a file
                line = f'{the_distance}>{result}'

                if line.endswith('and right over small jump 3-002'):
                    # ignore this line
                    continue

                lines_to_write.append({'distance':the_distance, 'text':result})
                # logger.info(f'recognized: {line}')

                # delete the tmp file
                # tmp_file.unlink()

        # sort lines_to_write
        lines_to_write = sorted(lines_to_write, key=lambda k: k['distance'])

        # write the result to a file
        if args['output_path']:
            output_path = args['output_path']
        elif args['tracks_path']:
            output_path = args['tracks_path']
        else:
            output_path = track_path
            
        output_file = Path(output_path) / f'{trackname.name}.pacenote'

        scriptGen = ScriptGenerator()
        with open(output_file, 'w') as f:
            for line in lines_to_write:
                x = f'{line["distance"]}>{line["text"]}'
                f.write(scriptGen.generate_pacenote_for_line(x))
                f.write('\n')

if __name__ == "__main__":
    main()

